package io.cici.cc.mybatis.lite.executor;

import io.cici.cc.mybatis.lite.executor.keygen.Jdbc3KeyGenerator;
import io.cici.cc.mybatis.lite.executor.keygen.KeyGenerator;
import io.cici.cc.mybatis.lite.executor.keygen.NoKeyGenerator;
import io.cici.cc.mybatis.lite.executor.resultset.ResultContextHandler;
import io.cici.cc.mybatis.lite.executor.statement.StatementHandler;
import io.cici.cc.mybatis.lite.mapping.MappedStatement;
import io.cici.cc.mybatis.lite.mapping.Sql;
import io.cici.cc.mybatis.lite.session.Configuration;
import io.cici.cc.mybatis.lite.transaction.Transaction;

import java.sql.BatchUpdateException;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

public class BatchExecutor extends BaseExecutor {

    public static final int BATCH_UPDATE_RETURN_VALUE = Integer.MIN_VALUE + 1002;

    private final List<Statement> statementList = new ArrayList<>();
    private final List<BatchResult> batchResultList = new ArrayList<>();
    private String currentSql;
    private MappedStatement currentStatement;

    public BatchExecutor(Configuration configuration, Transaction transaction) {
        super(configuration, transaction);
    }

    @Override
    public int doUpdate(MappedStatement ms, Object parameterObject) throws SQLException {
        final Configuration configuration = ms.getConfiguration();
        final StatementHandler statementHandler = configuration.newStatementHandler(this, ms, parameterObject, null,
                null);
        final Sql boundSql = statementHandler.getBoundSql();
        final String sql = boundSql.getSql();
        final Statement statement;
        if (sql.equals(currentSql) && ms.equals(currentStatement)) {
            int last = statementList.size() - 1;
            statement = statementList.get(last);

            statementHandler.parameterize(statement);// fix Issues 322
            BatchResult batchResult = batchResultList.get(last);
            batchResult.addParameterObject(parameterObject);
        } else {
            Connection connection = getConnection();
            statement = statementHandler.prepare(connection);
            statementHandler.parameterize(statement); // fix Issues 322
            currentSql = sql;
            currentStatement = ms;
            statementList.add(statement);
            batchResultList.add(new BatchResult(ms, sql, parameterObject));
        }
        statementHandler.batch(statement);
        return BATCH_UPDATE_RETURN_VALUE;
    }

    @Override
    public <E> List<E> doQuery(MappedStatement mappedStatement, Object parameterObject,
                               ResultContextHandler resultContextHandler, Sql sql) throws SQLException {
        Statement statement = null;
        try {
            flushStatements();
            Configuration configuration = mappedStatement.getConfiguration();
            StatementHandler statementHandler = configuration.newStatementHandler(wrapper, mappedStatement,
                    parameterObject, resultContextHandler, sql);
            Connection connection = getConnection();
            statement = statementHandler.prepare(connection);
            statementHandler.parameterize(statement);
            return statementHandler.query(statement, resultContextHandler);
        } finally {
            closeStatement(statement);
        }
    }

    @Override
    public List<BatchResult> doFlushStatements(boolean isRollback) throws SQLException {
        try {
            List<BatchResult> results = new ArrayList<>();
            if (isRollback) {
                return Collections.emptyList();
            }
            for (int i = 0, n = statementList.size(); i < n; i++) {
                Statement stmt = statementList.get(i);

                BatchResult batchResult = batchResultList.get(i);
                try {
                    batchResult.setUpdateCounts(stmt.executeBatch());
                    MappedStatement mappedStatement = batchResult.getMappedStatement();
                    List<Object> parameterObjects = batchResult.getParameterObjects();
                    KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
                    if (Jdbc3KeyGenerator.class.equals(keyGenerator.getClass())) {
                        Jdbc3KeyGenerator jdbc3KeyGenerator = (Jdbc3KeyGenerator) keyGenerator;
                        jdbc3KeyGenerator.processBatch(mappedStatement, stmt, parameterObjects);
                    } else if (!NoKeyGenerator.class.equals(keyGenerator.getClass())) { // issue #141
                        for (Object parameter : parameterObjects) {
                            keyGenerator.processAfter(this, mappedStatement, stmt, parameter);
                        }
                    }
                    // Close statement to close cursor #1109
                    closeStatement(stmt);
                } catch (BatchUpdateException e) {
                    StringBuilder message = new StringBuilder();
                    message.append(batchResult.getMappedStatement().getId()).append(" (batch index #").append(i + 1)
                            .append(")").append(" failed.");
                    if (i > 0) {
                        message.append(" ").append(i)
                                .append(" prior sub executor(s) completed successfully, but will be rolled back.");
                    }
                    throw new RuntimeException(message.toString(), e);
                }
                results.add(batchResult);
            }
            return results;
        } finally {
            for (Statement stmt : statementList) {
                closeStatement(stmt);
            }
            currentSql = null;
            statementList.clear();
            batchResultList.clear();
        }
    }

}
